package net.pragyah.scalby.dseparation

import scala.collection.mutable.Map

import net.pragyah.scalgorithms.graphs._


/* 
 * This implementation of DSeparation algorithm is taken from the book  
 * "Learning Bayesian Networks" by Richard E. Neapolitan
 * Looks like it is based on the algorithm by Dan Geiger, Thomas Verma & Judea Pearl
 * proposed in 
 * Proceedings of the Fifth Annual Conference on Uncertainty in Artificial Intelligence table of contents
 * Pages: 139 - 148  
 * Year of Publication: 1990
 * ISBN:0-444-88738-5
 */
class DSeparationAlgoGVP extends DSeparationAlgorithm{
  
  /*
   * Given a DAG G ...determine the set of all nodes d-separated from one
   * set of nodes,B by another set of nodes,A
   */
   
   def findDSeparations(G:Graph[char],
                       A:Set[Vertex[char]],
                       B:Set[Vertex[char]]) 
  : Set[Vertex[char]] = {
    assume(G.directed)

    val oldTags = Map[Vertex[char],Any]() 
    G.vertices.foreach(V => {
                        oldTags += V -> V.tag
                        V.tag = (false,false,false)
                       }) //keep the original tags aside ... assgn new tags .. for this subroutine
    
    //keep the old weights aside ... would need to set them back before returning
    val oldWeights= Map[Edge[char],double]() 
    G.edges.foreach(e => {oldWeights += e -> e.weight} )

    //TAG is made up of three booleans ... 
    //1st signifies whether the vertex is in A or not .. 
    //2nd signifies whether the vertex is in A or it has a decendent in A or not 
    // 3rd is used by the isOrHasDescInA routine to mark the visited nodes.
    
    G.vertices.foreach(V =>{
      val tag = V.tag.asInstanceOf[Tuple3[boolean,boolean,boolean]]
      //if this node has not been visited alredy ..... mark the relevant flags by visiting isOrHasDescInA routine
      if(!tag._3){
        isOrHasDescInA(V,A)
      }
      //tag associated with each vertex is a tuple consisting of the 'in' and 'descendent' flags
      //the third boolean is to find out if this node is already handled or not 
      
    })
    
    //Temporarily add the reverse-edges in the graph ....
    //maintain a set of all the temporarily added edges ... they'll have to be removed eventually
      // not just that .. this edges have to be discounted while finding out if a pat is legal or not (see isLegal)
    var exclusionSet = Set[Edge[char]]() 
    G.edges.foreach(e => {
      if(G.getEdge(e.v2,e.v1) == None){
        val edge = G.addEdge(e.v2,e.v1)
        exclusionSet = exclusionSet + edge
      }
    })
    

    
    val R = find_reachable_nodes(G,B,isLegal(G,exclusionSet) _)
    
    //D = V - (A U R) We do not need to remove B because B is a subset of R
    val AuR = R ++ A
    val D = Set()++G.vertices.filter(!AuR.contains(_))
    
    G.vertices.foreach(v => v.tag = oldTags.get(v)) //reset the tags .. to the original values
    exclusionSet.foreach(e => G.removeEdge(e.v1,e.v2))
    G.edges.foreach(e => {e.weight = oldWeights(e) } )

    D
  }


  /*
   * determines if a vertex is in A or any of its descendants are in A
   * Assumes that the vertices have tag of type Tuple3[boolean,boolean,boolean]
   * if the vertex has not been explored yet (i.e. tag._3 = false) 
   * it is determined if the vertex is in A and if it is then tag._1 is marked as true
   * it is also determined if the vertex or any of its Descendants is in A and tag._2 is marked accordingly.
   * 
   */
  def isOrHasDescInA(V:Vertex[char],A:Set[Vertex[char]]) : boolean = {
    var tag = V.tag.asInstanceOf[Tuple3[boolean,boolean,boolean]]
    if(tag._3) return tag._2 || tag._1 //if this tag has already been visited ... let kmow if it is or has a descendant in A
    // if V is a member of A .. then set flag1 to true .. flag2 to true .. and visited is true as well.
    if(A.contains(V)){
      V.tag = (true,true,true)// set - 'is' flag to true ... 'isOrHas' flag to true .. and visited to true 
      return true //yay!!! i'm in A!!
    } 
    
    V.tag = (false,false,true)//set the visited flag to true .. 'is' is false ... 'isOrHas' is not determined yet

    V.getEdges.foreach(e =>{
                         val o = e.getOther(V)
                         val oTag = o.tag.asInstanceOf[Tuple3[boolean,boolean,boolean]]

                         if(oTag._1 || oTag._2 ||  (!oTag._3 && isOrHasDescInA(o,A))){
                           //if 
                           // i) it is alreadyt determined that the vertex o exists in A
                            // OR ii) it is already determined that one of its descendent is in A
                             // OR iii) ... the vertex has not been visied && it is now determied that itself or one of its  decendent  is in A
                             //then return true
                              
                            V.tag = (false,true,true)//  'is' flag continues to be false... 'isOrHas' flag istrue .. and visited is true as well 
                           return true;//OK .. i may not be in A .. but i know of someone (directly or indirectly) who's in A
                         } 
                       } 
    )
    false // i'm not in A and i don't have any of my descendant's in A either :( 
  }
  
  
  def isLegal(G:Graph[char],exclusionSet:Set[Edge[char]])(U:Vertex[char],V:Vertex[char],W:Vertex[char]):boolean = {//TODO write test case for this one
      //The pair (U->V,V->W) is legal if and only if U != W and one of the following hold : 
      // i) U - V - W is not head-to-head in G and in[V] is false;
      // ii) U - V - W is head-to-head in G and descendent[v] is true.
      
      var UV = G.getEdge(U,V)
      if(UV != None && exclusionSet.contains(UV.get)) UV = None
      var WV = G.getEdge(W,V)
      if(WV != None && exclusionSet.contains(WV.get)) WV = None
      
      (U != W) &&
        // i)U - V - W is not head-to-head in G and in[V] is false;
        ((( UV == None ||  WV == None) 
          && !V.tag.asInstanceOf[Tuple3[boolean,boolean,boolean]]._1) // this means that V is not in A
         || //OR ii) U - V - W is head-to-head in G and descendent[v] is true.
         ((UV != None && WV != None) 
          && V.tag.asInstanceOf[Tuple3[boolean,boolean,boolean]]._2 // this means that V or its descendant is in A
         ))
    }    

  
  /*
   * Analysis 
   * No of nodes : n, No of edges : m 
   * Worst Case time complexity = O(mn)
   */
  
  
  
  def find_reachable_nodes(G:Graph[char],
                           B:Set[Vertex[char]], 
                           isLegal:(Vertex[char],Vertex[char],Vertex[char]) => boolean) : Set[Vertex[char]] = {
    assume(G.directed)
    
    var R = Set[Vertex[char]]()
    val iUV = Map[int,List[Tuple2[Vertex[char],Vertex[char]]]]()
    
    
    // get to all the children of each vertex in B and 
    // - mark the corresponding edges with a 1
    // - add the child node in the return set
    iUV += 1 -> List[Tuple2[Vertex[char],Vertex[char]]]()
    B.foreach(X => {
               R = R + X
               X.getEdges.foreach( e => {
                 val V = e.getOther(X)
                 R = R + V // add to the return set
                 iUV(1) = (X,V)::iUV(1) // add to the map for later use 
                 e.weight = 1 // freely changing the weight of the edges .. assuming they are of no use in this scenario  
                 })
              })

    var i = 1
    var found = true; // tells me if i have found nodes in the next level or not
    
    while(found){ // worst-case  - m times  (in between this while block and the iUV(i).foreach block inside it
      // this loop can go on for m cycles in the worst case, where m is the number of vertices
      
      found = false //assume for now that this is the last level ...
      iUV -= (i-1) // ah ... might as well clean it up  if there is anything .. you never know how big the graph is
      iUV += (i+1) -> List[Tuple2[Vertex[char],Vertex[char]]]()  // add the placeholder for the next level
      
      iUV(i).foreach(UV => { // for each V such that U->V is labeled i
        val U = UV._1
        val V = UV._2
        V.getEdges.filter(_.weight == 0).filter(e => isLegal(U,V,e.getOther(V))).foreach(e => {// worst case - n times .. for each vertex .... hence total cost of this method is O(mn) 
          // for each unlabelled edge V->W such that  (U -> V ,V -> W ) is legal 
          //.. in scalgorithms ...  0 weight is no weight
              val W = e.getOther(V)
              R = R+W
              e.weight = i+1
              iUV(i+1) = (V,W)::iUV(i+1) // add to the map for later use 
              found = true
           }
        )
      })
      i = i+1
    }
    R
  }
  
}
